"""
评估程序：实现针对每个测试集中的用户对召回的菜品和用户购买过的菜品进行统一打分，排序，最后对每个用户生成推荐的菜品（竞赛中最终是50个菜品）
"""
from .metrics import *
from .parser import parse_args

import torch
import numpy as np
import multiprocessing
import heapq
from collections import defaultdict

cores = multiprocessing.cpu_count() // 2

args = parse_args()
Ks = eval(args.Ks)
device = torch.device("cuda:" + str(args.gpu_id)) if args.cuda else torch.device("cpu")
BATCH_SIZE = args.test_batch_size
batch_test_flag = args.batch_test_flag

'''
针对每个用户排序模型输出的打分
'''
def ranklist_by_heapq(user_pos_test, test_items, rating, Ks, u):
    item_score = {}
    train_item_score = {}
    for i in test_items:
        item_score[i] = rating[i]
    if args.top_generate_setting == 'include_train':
        for i in train_user_set[u]:
            train_item_score[i] = rating[i]

    K_max = max(Ks)
    K_max_item_score = heapq.nlargest(K_max, item_score, key=item_score.get)
    if args.top_generate_setting == 'include_train':
        # 可以调节加入多少个用户之前买过的菜品
        K_max_train_item_score = heapq.nlargest(K_max, train_item_score, key=train_item_score.get)
        r = []
        count = 0
        top_result_items = []
        # top_result_items_scores = []
        for i in K_max_train_item_score:
            if count == args.rebuy_num:
                break
            top_result_items.append(i)
            if i in user_pos_test:
                r.append(1)
            else:
                r.append(0)
            count += 1
        for i in K_max_item_score:
            if count == K_max:
                break
            top_result_items.append(i)
            if i in user_pos_test:
                r.append(1)
            else:
                r.append(0)
            count += 1
        auc = 0.

        for i in top_result_items:
            top_result[u].append(i)
        return r, auc

    elif args.top_generate_setting == 'no_train':
        # if args.mode == 'evaluate':
        for i in K_max_item_score:
            top_result[u].append(i)

        r = []
        for i in K_max_item_score:
            if i in user_pos_test:
                r.append(1)
            else:
                r.append(0)
        auc = 0.
        return r, auc


def get_auc(item_score, user_pos_test):
    item_score = sorted(item_score.items(), key=lambda kv: kv[1])
    item_score.reverse()
    item_sort = [x[0] for x in item_score]
    posterior = [x[1] for x in item_score]

    r = []
    for i in item_sort:
        if i in user_pos_test:
            r.append(1)
        else:
            r.append(0)
    auc = AUC(ground_truth=r, prediction=posterior)
    return auc


def ranklist_by_sorted(user_pos_test, test_items, rating, Ks):
    item_score = {}
    for i in test_items:
        item_score[i] = rating[i]

    K_max = max(Ks)
    K_max_item_score = heapq.nlargest(K_max, item_score, key=item_score.get)

    r = []
    for i in K_max_item_score:
        if i in user_pos_test:
            r.append(1)
        else:
            r.append(0)
    auc = get_auc(item_score, user_pos_test)
    return r, auc


def get_performance(user_pos_test, r, auc, Ks):
    hit_ratio = []

    for K in Ks:
        hit_ratio.append(hit_at_k(r, K))

    return {'hit_ratio': np.array(hit_ratio)}


def test_one_user(x):
    # user u's ratings for user u
    rating = x[0]
    # uid
    u = x[1]
    try:
        training_items = train_user_set[u]
    except Exception:
        training_items = []
    # user u's items in the test set
    user_pos_test = test_user_set[u]

    test_items = list(set(recall_items) - set(training_items))

    if args.test_flag == 'part':
        r, auc = ranklist_by_heapq(user_pos_test, test_items, rating, Ks, u)
    else:
        r, auc = ranklist_by_sorted(user_pos_test, test_items, rating, Ks)

    return get_performance(user_pos_test, r, auc, Ks)


def test(model, user_dict, n_params, recall_item, user_rebuys):
    result = {
        'hit_ratio': np.zeros(len(Ks)),
    }
    global top_result, top_result_score
    top_result = defaultdict(list)
    global n_users, n_items, recall_items, user_rebuy
    n_items = n_params['n_items']
    n_users = n_params['n_users']
    recall_items = recall_item
    user_rebuy = user_rebuys

    global train_user_set, test_user_set
    train_user_set = user_dict['train_user_set']
    test_user_set = user_dict['test_user_set']

    pool = multiprocessing.Pool(cores)

    u_batch_size = BATCH_SIZE
    i_batch_size = BATCH_SIZE

    test_users = list(test_user_set.keys())
    n_test_users = len(test_users)
    n_user_batchs = n_test_users // u_batch_size + 1

    count = 0

    entity_gcn_emb, user_gcn_emb = model.generate()  # 这个位置是不是返回了三个但是只接受了两个？是的，注意这个上面有一个[:-1]

    for u_batch_id in range(n_user_batchs):
        start = u_batch_id * u_batch_size
        end = (u_batch_id + 1) * u_batch_size

        user_list_batch = test_users[start: end]
        user_batch = torch.LongTensor(np.array(user_list_batch)).to(device)
        u_g_embeddings = user_gcn_emb[user_batch]

        if batch_test_flag:
            n_item_batchs = len(recall_items) // i_batch_size + 1
            rate_batch = np.zeros(shape=(len(user_batch), n_items))

            i_count = 0
            for i_batch_id in range(n_item_batchs):
                i_start = i_batch_id * i_batch_size
                i_end = min((i_batch_id + 1) * i_batch_size, len(recall_items))

                item_batch = torch.LongTensor(recall_items[i_start: i_end]).view(i_end - i_start).to(device)
                i_g_embddings = entity_gcn_emb[item_batch]

                i_rate_batch = model.rating(u_g_embeddings, i_g_embddings).detach().cpu()

                rate_batch[:, recall_items[i_start: i_end]] = i_rate_batch
                i_count += i_rate_batch.shape[1]

            assert i_count == len(recall_items)
        else:
            item_batch = torch.LongTensor(recall_items).to(device)
            i_g_embddings = entity_gcn_emb[item_batch]
            rate_batch = model.rating(u_g_embeddings, i_g_embddings).detach().cpu()

        user_batch_rating_uid = zip(rate_batch, user_list_batch)
        if False:
            batch_result = pool.map(test_one_user, user_batch_rating_uid)
        else:  # 对于windows系统只能先这样，不能用这个程序的多线程。具体原因还不清楚。
            batch_result = []
            for x in user_batch_rating_uid:
                batch_result.append(test_one_user(x))
        count += len(batch_result)

        for re in batch_result:
            result['hit_ratio'] += re['hit_ratio'] / n_test_users

    assert count == n_test_users
    pool.close()

    # 把每个用户推荐的top-n物品写入文件
    with open('result/' + str(args.save_model_id) + '_top_' + str(args.rebuy_num) + '_' + str(args.dataset) + '_' + str(
            args.model) + '_' + str(args.test_set) + '.txt', 'w') as f:
        for key in top_result:
            f.write(str(key))
            for i in top_result[key]:
                f.write('\t')
                f.write(str(i))
            f.write('\n')

    return result
